from math import cos, pi, sin
from numpy import asarray, ndarray
from numpy.random import randint, random
from pygame import init, QUIT, Surface
from pygame.display import flip, set_mode, set_caption
from pygame.draw import line, circle
from pygame.event import get
from pygame.image import save

__all__ = [
    "genTree",
    "drawTree",
]


# 需要 numpy 和 pygame 这两个库

# Setting
LEVEL = 3                       # 树的生长阶段
SIZE = 1000, 1000               # 图片的长宽高
FILENAME = f"tree{LEVEL}.png"   # 生成文件名称
ITER = 7+LEVEL                  # 分叉迭代次数
LENG = 80.+LEVEL*30             # 枝条初始长度
ANGLE = .04*LEVEL               # 枝条散开角度
GREEN = 1+LEVEL*10              # 叶子颜色深度
IS_BLANK = False                # 背景是否透明


def genTree(health: int, length: float, root: ndarray, heading: float) -> "tuple[tuple[ndarray,ndarray,int],tuple[ndarray]]":
    """递归计算每一个枝条、叶子的位置和大小"""

    delta = asarray([cos(heading), sin(heading)])
    now = root+delta*length

    if health:
        lines, circles = ((root, now, health+1),), ()

        b = random()/1.5+ANGLE
        d = length*(random()/5+.55)
        line, circle = genTree(health-1, d, now, heading+b)
        lines += line
        circles += circle

        b = random()/1.5+ANGLE
        d = length*(random()/5+.55)
        line, circle = genTree(health-1, d, now, heading-b)
        lines += line
        circles += circle

        d = length*(random()/5+.65)
        line, circle = genTree(health-1, d, now, heading)
        lines += line
        circles += circle

        return lines, circles
    else:
        return ((root, now, health+1),), (now,)


def drawTree(surf: Surface, health: int, length: float) -> None:
    """把计算得到的树画出来"""
    root = asarray([surf.get_width() >> 1, surf.get_height()])
    lines, circles = genTree(health, length, root, -pi/2)
    color1 = asarray([128, 42, 42, 255])
    color2 = asarray([30, 255, 30, 255])
    mask1, mask2 = asarray([1, 1, 1, 0]), asarray([0, -1, 0, 0])
    offset = asarray([GREEN]*3+[0])

    for fr, to, width in lines:
        if width > 1:
            color = color1+randint(30, size=4)*mask1-offset
        else:
            color = color2+randint(GREEN, size=4)*mask2
        line(surf, color, fr, to, width)

    for pos in circles:
        offset1 = randint(GREEN*6, size=4)*mask2
        circle(surf, color2+offset1, pos, randint(1, 6))


if __name__ == "__main__":
    init()
    screen = set_mode(SIZE)
    set_caption("Print tree")
    screen.fill([255]*3)
    surf = Surface(SIZE).convert_alpha()
    surf.fill([0]*4)
    drawTree(surf, ITER, LENG)
    screen.blit(surf, screen.get_bounding_rect())
    flip()
    save(surf if IS_BLANK else screen, FILENAME)

    loop = 1
    while loop:
        for even in get():
            if even.type == QUIT:
                loop = 0
                break
